package com.shujia.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo02PersonTrain {
  def main(args: Array[String]): Unit = {
    // 创建SparkSession
    val spark: SparkSession = SparkSession
      .builder()
      .appName("Demo02PersonTrain")
      .master("local[*]")
      .config("spark.sql.shuffle.partitions", "8")
      .getOrCreate()

    // 1、读取数据 并对数据做特征工程处理
    val personDF: DataFrame = spark
      .read
      .format("libsvm") // 经过特征工程处理过后的数据
      .load("Spark/data/mllib/data/人体指标.txt")

    personDF.show(truncate = false)

    // 2、切分数据 形成 训练集和测试集 比例约为 8:2
    val array: Array[Dataset[Row]] = personDF
      .randomSplit(Array(0.8, 0.2))

    val trainDF: Dataset[Row] = array(0) // 训练集
    val testDF: Dataset[Row] = array(1) // 测试集

    /**
     * 3、选择合适的模型
     * 怎么选？
     * 有监督学习or无监督学习 --> 原数据有无label
     *    - 有label：有监督学习 --> 看label是连续的 还是 离散的
     *                            - 连续的：回归
     *                            - 离散的：分类 --> 选择逻辑回归模型
     *    - 无label：无监督学习
     */

    val logisticRegression: LogisticRegression = new LogisticRegression()
      .setMaxIter(10) // 最大的迭代次数
      .setRegParam(0.3)
      .setElasticNetParam(0.8)

    // 4、使用训练集进行训练 得到训练好的模型
    val logisticRegressionModel: LogisticRegressionModel = logisticRegression.fit(trainDF)

    // 5、使用测试集进行模型评估
    val testTransDF: DataFrame = logisticRegressionModel.transform(testDF)

    testTransDF.show(truncate = false)

    import spark.implicits._
    import org.apache.spark.sql.functions._

    /**
     * 准确率、精确率、召回率：
     * https://cloud.tencent.com/developer/article/1510724
     */

    // 计算模型的准确率：预测正确的/总的数据条数
    testTransDF
      .withColumn("isEqual", when($"label" === $"prediction", 1).otherwise(0))
      .select(sum($"isEqual") / count("*") as "准确率")
      .show()

    // 计算模型的精确率：数据中原来label为1预测结果也为1的数据条数/数据中预测结果为1（正性）的条数
    testTransDF
      .withColumn("TP", when($"label" === 1 and $"prediction" === 1, 1).otherwise(0))
      .withColumn("TP_FP", when($"prediction" === 1, 1).otherwise(0))
      .select(sum($"TP")/sum($"TP_FP") as "精确率")
      .show()

    // 6、模型评估通过后 可以将模型保存并使用
    logisticRegressionModel.save("Spark/data/mllib/person")


  }

}
